// 测试算子的速度，在各个设备上的表现。
#include "DpaUnit.h"
#include "timer.h"
#include "DeviceShare.h"

#define DEFAULT_WINDOWS_SIZE 1024 * 1024 * 1024

DEFINE_int32 (thread_num,  1, "thread for device to use");
DEFINE_string(rdma_device, "mlx5_bond_0", "rdma_device_name");
DEFINE_string(dpa_type, "flexio", "dpa sdk type"); 
DEFINE_int32(batch_size, 1, "batch_size");
DEFINE_int32(execute_time, 1, "execute times");

/* compare Dpa result with cpu's sigrid_hash */
void correctness_check(int64_t * input_ptr, int64_t * output_ptr, size_t size) {
    int incorrect = 0;
    for (int i = 0; i < (int)size; ++i) {
        int64_t x = input_ptr[i];
        uint64_t value = (x ^ (uint64_t)(0));
        value = (value * 2654435761U) ^ (value >> 16);
        value = value % (uint64_t)(10);

        int64_t y = output_ptr[i];
        if (value != (uint64_t)y) incorrect ++; 
    }
    printf ("incorrect position num = %d \n", incorrect);
}

int main(int argc, char **argv) {
    gflags::ParseCommandLineFlags(&argc, &argv, true);

    xmh::Reporter::StartReportThread();

    void * ptr = nullptr;
    int ret = posix_memalign(&ptr, 64, ( (size_t)DEFAULT_WINDOWS_SIZE + 63) & (~63));
    if (ret != 0) {
        printf ("error malloc host memory\n");
    }

    size_t size_input = (size_t)FLAGS_batch_size * (size_t)sizeof(double);

    void * input_mem_ptr  = ptr;
    void * output_mem_ptr = ptr + size_input;


    /* init input_mem_addr */
    printf ("input info: ");
    for (int i = 0; i < FLAGS_batch_size; ++i) {
        ( (int64_t *) (input_mem_ptr + i * sizeof(double))) [0] = 12; 
        if (i < 10) { printf ("[%ld] ", ( (int64_t *) (input_mem_ptr + i * sizeof(double))) [0]);}
    } printf ("\n");
    /* init input_mem_addr */

    DpaUnit * dpaunit;

    if(FLAGS_dpa_type == "docasdk") {
        dpaunit = new DpaDocasdkUnit(); 
    }
    else {
        dpaunit = new DpaFlexioUnit();
    }
    
    /* register mem window so dpa can access arm memory */
    dpaunit->register_host_memory(ptr, (size_t)DEFAULT_WINDOWS_SIZE );
    /* register mem window so dpa can access arm memory */

    /* Assign tasks to each thread in advance */
    DpaThreadInfo task_info;
    task_info.input_mem_addr  = input_mem_ptr;
    task_info.output_mem_addr = output_mem_ptr;
    task_info.thread_num      = FLAGS_thread_num;
    task_info.size            = FLAGS_batch_size;

    dpaunit->init_dpa_thread_info(task_info);
    /* Assign tasks to each thread in advance */
    
    /* task execute */
    for (int i = 0; i < FLAGS_execute_time; ++i) {
        // time calc 
        xmh::Timer dpa_timer(FLAGS_dpa_type); 

        // submit task to each thread 
        dpaunit->submit_task(0, FLAGS_thread_num);

        // wait until all threads finish there own task 
        dpaunit->poll_thread();
        dpa_timer.end();
    }

    /* simple check */
    printf ("output info: ");
    for (int i = 0; i < FLAGS_batch_size; ++i) {
        if (i > 10) break;
        printf ("[%ld] ", ( (int64_t *) (output_mem_ptr + i * sizeof(double)) ) [0]);
    } printf ("\n");
    /* simple check */

    /* carefully check */
    correctness_check( (int64_t *) input_mem_ptr, (int64_t *) output_mem_ptr, FLAGS_batch_size);
    /* carefully check */

    xmh::Reporter::StopReportThread();
    xmh::Reporter::Report();

    /* get max thread execute time */
    size_t max_execute_time_per_thread = dpaunit->get_execute_cycle();
    /* get max thread execute time */

    double max_execute_time_us = (double)1.0 * max_execute_time_per_thread * 1000000.00 / 1800000000.00;
    double total_execute_time_us = xmh::Timer::ManualQuery(FLAGS_dpa_type) / 1000.00;

    std::cout << FLAGS_dpa_type << " task_execute info" << std::endl;
    printf ("max_execute_time_per_thread = %.6lf us\n", max_execute_time_us);
    printf ("total_execute_time = %.6lf us\n", total_execute_time_us);

    free(ptr);

    sleep(2);

    return 0;
}

// ./../bin_dpu/test_operator --op_type="sigrid_hash" --thread_num=128 --device_type="dpa" --batch_size=8192 --dpa_type="0"